'''
Author: Q.C.Li
Date: 2020-10-21 09:14:54
LastEditors: Please set LastEditors
LastEditTime: 2022-09-30 09:02:09
Description: 一些绘图函数
'''
import webbrowser

import numpy as np
import plotly.express as px
import plotly.graph_objects as go
import plotly.io as pio
from plotly.subplots import make_subplots


pio.templates.default = "presentation"  #"presentation"
pio.renderers.default ='chrome'
webbrowser.register('chrome', None,webbrowser.BackgroundBrowser(r"C:\Program Files\Google\Chrome\Application\chrome.exe"))

# %% 通用绘图函数
def slidePlot(data):
    # 滑块绘图
    # data:二维绘图图数据，0-axis 为滑块的每一步，1-axis为数据
    fig = go.Figure()
    # 为滑块每一步添加数据
    for i in range(data.shape[0]):
        fig.add_trace(go.Scatter(y=abs(data[i,:]), visible=False))
    steps = []
    # 定义每一步的显示内容
    for i in range(40):
        steps.append({"method":"update",'args': [{'visible': 40*[False] }] }.copy())
        steps[i]['args'][0]['visible'][i] = True
    # 更新布局，将 steps 添加到布局中
    fig.update_layout(sliders=[{'steps':steps}])
    return fig

def Plot_3D(x,y,z):
    # 3D 曲面图
    fig = go.Figure(go.Surface(
        # 添加轮廓线
        contours={
            'x':{'show':True,'start':x[0],'end':x[-1],'size':(x[-1]-x[0])/10},
            'y':{'show':True,'start':y[0],'end':y[-1],'size':(y[-1]-y[0])/10}},
        # 添加数据
        z=z,x=x,y=x,opacity=.9,
        colorscale='jet',showscale=False))
    # 设置布局
    layout = {'template': 'plotly_white',
        'scene':{
            'zaxis':{'showticklabels':False,    # 隐藏坐标轴数字
                'showgrid': False,      # 隐藏网格
                'visible': False,       # 隐藏坐标轴
                'range':[-.5,1.5],}     # 设置坐标轴范围
                }}
    fig.update_layout(layout)
    return fig

def showColors():
    # 连续颜色图
    fig = px.colors.sequential.swatches()
    fig.show()
    # 离散颜色图
    fig = px.colors.diverging.swatches().update_layout(margin_b=10)
    fig.show()
    # 周期颜色图
    fig = px.colors.cyclical.swatches_cyclical()
    fig.show()
    fig = px.colors.cyclical.swatches()
    fig.show()

def imshow(z,x=None,y=None,labels={},flnm=None)->go.Figure:
    '''二维绘图

    Parameters
    ----------
    z : 2D array
    x : 1D array,optional
    y : 1D array,optional
    labels : list[xtitle,y title]
    flnm : string,optional
        filename to save html file, by default None (not save)
    Returns
    -------
    plotly...Figure
    '''
    fig = px.imshow(
        img=z,x=x,y=y,
        labels={'x':labels[0],'y':labels[1]}, )
    fig.update_yaxes(autorange=True)
    if flnm is not None:
        fig.write_html(flnm,include_mathjax = 'cdn')
    return fig

def imshowgl(z,x=None,y=None,colorscale='viridis'):
    # 使用 WebGl 渲染的 hetamap 图，可以调用显卡，交互更丝滑,但渲染会失真
    fig = go.Figure(data=go.Heatmapgl(
        z=z,x=x,y=y,colorscale=colorscale))
    return fig

def plot(x,y,names,labels):
    '''多条 plot 绘图

    Parameters
    ----------
    x : list[n]
        横坐标数据，列表等可枚举的类型
    y : list[n]
        横坐标数据，列表等可枚举的类型
    names : list[name]
        每条线的名字
    labels: dict
        一些组件的 title, 如坐标轴

    Returns
    -------
        plot...Figure
    '''
    fig = go.Figure()
    for k,v in enumerate(names):
        fig.add_traces(go.Scatter(
            x=x[k],y=y[k],name=v, ))
    fig.update_xaxes(title_text=labels['x'])
    fig.update_yaxes(title_text=labels['y'])
    return fig


plot(
    x=[[1,2,3],[3,5,7]], y=[[1,2,3],[3,5,7]],
    names=['恢复互强度', '外差功率谱'],
    labels={
        'x': 'x',
        'y': 'y',
})

# %% 模式分解绘图
def showModes(En,x,y):
    # 绘制多模式的 subplot
    i = 0
    sm,sn,_,_ = En.shape
    subplot_titles = [None]*sm*sn
    fig = make_subplots(
        rows=sm, cols=sn,
        subplot_titles=['untitled']*sm*sn);
    for key1 in range(sm):
        for key2 in range(sn):
            fig.add_trace(go.Heatmap(
                z=abs(En[key1,key2]),
                x=x, y=y,),
                row=1+key1, col=1+key2);
            fig.layout.annotations[i]['text'] = f'{key1}{key2}模';
            i += 1
    return fig

def showSDC(demo,name,x):
    # 显示谱相干度
    W = demo.cmCsdSdSdc(name,0)
    fig1 = px.imshow(abs(W['SDC']),x=x,y=x,
        title=f'谱相干度\n{name}',
        labels={'x':r'$x_1$','y':r'$x_2$'})
    muDelta = abs(W['SDC'][::-1]).diagonal()
    fig2 = px.line(y=muDelta,x=2*x,
        title=r'$谱相干度\quad vs\quad\Delta x$',
        labels={'x':r'$\Delta x$','y':r'$\mu$'})
    return fig1,fig2

def CD_I(CD,I,x,title):
    fig1 = px.line(y=abs(I),x=x,
        labels={'x':r'x','y':'I'})
    fig2 = px.line(y=abs(CD),x=2*x,
        labels={'x':r'$\Delta x$','y':r'$\mu$'})
    fig_subplot = make_subplots(
        rows=2, cols=1,
        subplot_titles=[f'光强',r'$谱相干度\quad vs\quad\Delta x$']);
    fig_subplot.add_trace(
        fig1.data[0],
        row=1,col=1)
    fig_subplot.add_trace(
        fig2.data[0],
        row=2,col=1)
    fig_subplot.update_layout(title_text=title)
    fig_subplot.show()
    return fig_subplot

def intensity(demo, name):
    # 光强
    beta_m = demo.Wf[name]['beta-mn'][0].reshape(-1, 1, 1, 1)
    beta_n = demo.Wf[name]['beta-mn'][1].reshape(1, -1, 1, 1)
    return np.sum(abs(demo.Wf[name]['mode'])**2 * (beta_m*beta_n), axis=(0, 1))


# %% SRW 绘图




if __name__ == '__main__':
    # subplot 绘图
    fig = make_subplots(
        rows=2, cols=2,
        subplot_titles=['title']*4)
    fig.add_trace(go.Scatter(
        x=[1, 2, 3], y=[4, 5, 6],name='trace-lqc'),
                row=1, col=1);
    fig.add_trace(go.Scatter(x=[20, 30, 40], y=[50, 60, 70]),
                row=1, col=2);
    fig.add_trace(go.Scatter(x=[300, 400, 500], y=[600, 700, 800]),
                row=2, col=1);
    fig.add_trace(go.Scatter(
        x=[4000, 5000, 6000], y=[7000, 8000, 9000]),row=2, col=2);
    fig.update_layout(height=500, width=700,
                    title_text="Multiple Subplots with Titles");
    # for loop 中更改 subplot 的 title
    for ii in range(4):
        fig.layout.annotations[ii]['text'] = f'title{ii}'
    fig.show()

    # 二维影像图
    fig = go.Figure()
    fig.add_trace(go.Heatmap(z=np.random.poisson(size=[100,100]),
        x=np.linspace(-10,10,100),
        y=np.linspace(-1,3,100),
        colorscale='hot',showscale=False ))
    px.imshow(np.random.poisson(size=[100,100]),color_continuous_scale ='hot')

    # 坐标轴反转
    fig.update_layout(yaxis_autorange=True);
    
